# 导入函数库
from jqdata import *
import math

enable_profile()


# 初始化函数，设定基准等等
def initialize(context):
    # 设定沪深300作为基准
    set_benchmark('000300.XSHG')
    # 开启动态复权模式(真实价格)
    set_option('use_real_price', True)
    # 输出内容到日志 log.info()
    log.info('初始函数开始运行且全局只运行一次')
    # 过滤掉order系列API产生的比error级别低的log
    log.set_level('order', 'error')

    ### 股票相关设定 ###
    # 股票类每笔交易时的手续费是：买入时佣金万分之三，卖出时佣金万分之三加千分之一印花税, 每笔交易佣金最低扣5块钱
    set_order_cost(OrderCost(close_tax=0.001, open_commission=0.00015, close_commission=0.00015, min_commission=5),
                   type='stock')

    ## 运行函数（reference_security为运行时间的参考标的；传入的标的只做种类区分，因此传入'000300.XSHG'或'510300.XSHG'是一样的）

    # 开盘时运行
    run_daily(trade, time='09:28')


def trade(context):
    # 先卖出股票腾仓位
    # 看看持有股票，如果有低于5日线的，卖出
    long_positions_dict = context.portfolio.long_positions
    for position in list(long_positions_dict.values()):
        # 9:30之前，position.price 还是昨天的收盘价
        # 取对应的集合竞价后的价格
        current_price = get_call_auction(position.security, start_date=context.current_dt, end_date=context.current_dt,
                                         fields=['current'])['current'][0]
        # 止盈（50%） 止损（10%）
        if current_price < position.acc_avg_cost * 0.9:
            order_target(position.security, 0)
            log.info("止损（10%），卖出=======" + position.security)
        if current_price > position.acc_avg_cost * 1.5:
            order_target(position.security, 0)
            log.info("止盈（50%），卖出=======" + position.security)

        ma5 = get_bars(position.security, count=5, unit='1d', fields=['close'])['close'].mean()
        if current_price < ma5:
            order_target(position.security, 0)
            log.info("跌破5日线，卖出=======" + position.security)

        # 开盘价跌幅大于2个点，卖
        pre_close = attribute_history(position.security, 1, '1d', ('close'))['close'][-1]
        if current_price < pre_close * 0.98:
            order_target(position.security, 0)
            log.info("开盘跌幅大于2个点，卖出=======" + position.security)

    # 如果没有仓位了，今天不动了
    available_cash = context.portfolio.available_cash
    if available_cash < 3000:
        log.info("没钱了，今天不动了")
        return

    # 连续3天涨停的票
    stocks = high_limit_3d_stock(context)
    log.info("连续3天涨停的票=======" + str(stocks))
    # 账户总资产（现金+股票市值）, 每次买入一只票最多5成仓
    buy_max_per_stock = round(context.portfolio.total_value * 0.5, 2)
    # 循环处理每只备选票
    for stock in stocks:
        # 取得当前的现金
        cash = context.portfolio.available_cash
        if cash < 3000:
            log.info("没钱了，今天不继续买了")
            return
        # 昨天的收盘价
        pre_close = attribute_history(stock, 1, '1d', ('close'))['close'][-1]
        # 集合竞价价格
        current_price = \
        get_call_auction(stock, start_date=context.current_dt, end_date=context.current_dt, fields=['current'])[
            'current'][0]
        # 涨停价, 区分创业板、st，科创板
        high_limit_price = pre_close * 1.1
        if stock.startswith('3') or stock.startswith('68'):
            high_limit_price = pre_close * 1.2
        if stock.startswith('ST'):
            high_limit_price = pre_close * 1.05

        # 保留2位小数
        high_limit_price = round(high_limit_price, 2)

        # 跌停价
        # low_limit_price = current_data[stock].low_limit
        # 涨跌幅
        change_percent = round(current_price / pre_close - 1, 2)
        log.info("开盘价%f, 涨跌幅%f, =======%s" % (current_price, change_percent, stock))
        # 涨2个点以上，没涨停，买入
        if (current_price < high_limit_price) and (current_price > (pre_close * 1.02)) and (
                context.portfolio.positions[stock].closeable_amount <= 0):
            # 用所有 cash 买入股票
            # order_value(stock, cash)
            buy_cash = min(cash, buy_max_per_stock)
            order_value(stock, buy_cash)
            log.info("开盘价%f, 涨跌幅%f, 涨停价%f, 买入=======%s" % (
            current_price, change_percent, high_limit_price, stock))


def high_limit_3d_stock(context):
    # 获取所有上市股票
    all_stock = list(get_all_securities(['stock'], date=context.current_dt).index)
    # 获取连续3天涨停的票
    high_limit_3d = []
    for stock in all_stock:
        stock_history = attribute_history(stock, 3, '1d', ('high_limit', 'close'), df=True)
        flag = True
        for index, row in stock_history.iterrows():
            if (math.isnan(row['close'])) or (row['close'] != row['high_limit']):
                flag = False
                break
        if flag:
            high_limit_3d.append(stock)

    # 返回股票代码列表
    return high_limit_3d


def get_high_limit_stocks:









